from keras import backend as K
from keras.layers import *
from keras.layers.convolutional import *
from keras.layers.core import *
from keras.layers.normalization import BatchNormalization
from keras.layers.pooling import *
from keras.models import Model
from keras.regularizers import l2


class VggModelBuilder(object):
    @staticmethod
    def build(num_outputs=(5, 4, 6, 5, 5, 5, 5, 5, 5), name_outputs=None, weight_decay=3e-4):
        input = Input(shape=(64, 64, 64, 1))
        # block1
        x = Convolution3D(nb_filter=32, kernel_dim1=3, kernel_dim2=3, kernel_dim3=3,
                          init="he_normal", border_mode="same", W_regularizer=l2(weight_decay))(input)
        x = BatchNormalization()(x)
        x = Activation('relu')(x)
        x = Convolution3D(nb_filter=32, kernel_dim1=3, kernel_dim2=3, kernel_dim3=3,
                          init="he_normal", border_mode="same", W_regularizer=l2(weight_decay))(x)
        x = BatchNormalization()(x)
        x = Activation('relu')(x)
        x = MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2), border_mode='same')(x)

        # block2
        x = Convolution3D(nb_filter=64, kernel_dim1=3, kernel_dim2=3, kernel_dim3=3,
                          init="he_normal", border_mode="same", W_regularizer=l2(weight_decay))(x)
        x = BatchNormalization()(x)
        x = Activation('relu')(x)
        x = Convolution3D(nb_filter=64, kernel_dim1=3, kernel_dim2=3, kernel_dim3=3,
                          init="he_normal", border_mode="same", W_regularizer=l2(weight_decay))(x)
        x = BatchNormalization()(x)
        x = Activation('relu')(x)
        x = MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2), border_mode='same')(x)

        # block3
        x = Convolution3D(nb_filter=128, kernel_dim1=3, kernel_dim2=3, kernel_dim3=3,
                          init="he_normal", border_mode="same", W_regularizer=l2(weight_decay))(x)
        x = BatchNormalization()(x)
        x = Activation('relu')(x)
        x = Convolution3D(nb_filter=128, kernel_dim1=3, kernel_dim2=3, kernel_dim3=3,
                          init="he_normal", border_mode="same", W_regularizer=l2(weight_decay))(x)
        x = BatchNormalization()(x)
        x = Activation('relu')(x)
        x = Convolution3D(nb_filter=128, kernel_dim1=3, kernel_dim2=3, kernel_dim3=3,
                          init="he_normal", border_mode="same", W_regularizer=l2(weight_decay))(x)
        x = BatchNormalization()(x)
        x = Activation('relu')(x)
        x = MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2), border_mode='same')(x)

        # block4
        x = Convolution3D(nb_filter=256, kernel_dim1=3, kernel_dim2=3, kernel_dim3=3,
                          init="he_normal", border_mode="same", W_regularizer=l2(weight_decay))(x)
        x = BatchNormalization()(x)
        x = Activation('relu')(x)
        x = Convolution3D(nb_filter=256, kernel_dim1=3, kernel_dim2=3, kernel_dim3=3,
                          init="he_normal", border_mode="same", W_regularizer=l2(weight_decay))(x)
        x = BatchNormalization()(x)
        x = Activation('relu')(x)
        x = Convolution3D(nb_filter=256, kernel_dim1=3, kernel_dim2=3, kernel_dim3=3,
                          init="he_normal", border_mode="same", W_regularizer=l2(weight_decay))(x)
        x = BatchNormalization()(x)
        x = Activation('relu')(x)
        x = MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2), border_mode='same')(x)

        # block4
        flatten = Flatten()(x)
        dense1 = Dense(output_dim=512, init='he_normal', activation='relu', W_regularizer=l2(weight_decay))(flatten)
        dense3 = []
        for i, num_output in enumerate(num_outputs):
            if name_outputs is not None:
                dense2 = Dense(output_dim=256, init='he_normal', activation='relu', W_regularizer=l2(weight_decay)
                               )(dense1)
                dense3.append(
                    Dense(output_dim=num_output, init='he_normal', activation='softmax', W_regularizer=l2(weight_decay),
                          name=name_outputs[i])(dense2))
            else:
                dense2 = Dense(output_dim=256, init='he_normal', activation='relu', W_regularizer=l2(weight_decay))(
                    dense1)
                dense3.append(Dense(output_dim=num_output, init='he_normal', activation='softmax',
                                    W_regularizer=l2(weight_decay))(dense2))

        model = Model(input=input, output=dense3)
        return model


def main():
    model = VggModelBuilder.build(name_outputs=('Sub', 'IS', 'Cal', 'Sph', 'Mar', 'Lob', 'Spi', 'Tex', 'Mal'))
    model.compile(loss=["categorical_crossentropy"] * 9, optimizer="sgd")
    model.summary()


if __name__ == '__main__':
    main()
